import os
import matplotlib.pyplot as plt

# 定义一个 Plotter 类
class Plotter:
    # 初始化方法，传入参数包括两个日志目录、输出目录、两个文件名、曲线名称、坐标轴标签、图例位置和输出图片名
    def __init__(self, log_dir_A, log_dir_B, output_dir, file_A, file_B, A_curve, B_curve, X_label, Y_label, legend, output_png):
        self.log_dir_A = log_dir_A  # 日志目录 A
        self.log_dir_B = log_dir_B  # 日志目录 B
        self.output_dir = output_dir  # 输出目录
        self.file_A = os.path.join(log_dir_A, file_A)  # 拼接日志目录 A 和文件名 A，得到完整路径
        self.file_B = os.path.join(log_dir_B, file_B)  # 拼接日志目录 B 和文件名 B，得到完整路径
        self.A_curve = A_curve  # 曲线 A 的名称
        self.B_curve = B_curve  # 曲线 B 的名称
        self.X_label = X_label  # X 轴标签
        self.Y_label = Y_label  # Y 轴标签
        self.legend = legend  # 图例位置
        self.output_png = output_png  # 输出图片的文件名
        self.A = self.load(self.file_A)  # 加载文件 A 的数据
        self.B = self.load(self.file_B)  # 加载文件 B 的数据

    # 定义加载文件数据的方法
    def load(self, file_path):
        try:
            with open(file_path, 'r') as file:  # 打开文件
                # 读取文件中的每一行，并转换为浮点数，存储在列表中
                num = [float(line.strip()) for line in file.readlines()]
            return num  # 返回数据列表
        except FileNotFoundError:  # 如果文件未找到
            return []  # 返回空列表

    # 定义绘图方法
    def plot(self):
        iters_A = range(len(self.A))  # A 数据的迭代次数
        iters_B = range(len(self.B))  # B 数据的迭代次数

        plt.figure()  # 创建一个新图形
        plt.plot(iters_A, self.A, 'red', linewidth=2, label=self.A_curve)  # 绘制 A 曲线
        plt.plot(iters_B, self.B, 'blue', linewidth=2, label=self.B_curve)  # 绘制 B 曲线
        plt.grid(True)  # 显示网格
        plt.xlabel(self.X_label)  # 设置 X 轴标签
        plt.ylabel(self.Y_label)  # 设置 Y 轴标签
        plt.legend(loc=self.legend)  # 设置图例位置

        # 如果输出目录不存在，则创建
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

        # 保存图形到指定的文件
        plt.savefig(os.path.join(self.output_dir, self.output_png))
        plt.cla()  # 清除当前活动的图形
        plt.close("all")  # 关闭所有图形窗口

# 如果此脚本是主程序，则执行以下代码
if __name__ == "__main__":
    # 设置日志目录和输出目录
    log_directory_A = r"E:\postgraduate\大论文\seg_train\短数据集 500轮次"
    log_directory_B = r"E:\postgraduate\大论文\seg_train\短数据集 500轮次"
    output_directory_C = r"E:\postgraduate\大论文\seg_train\短数据集 500轮次"

    # 设置文件名、曲线名称、坐标轴标签、图例位置和输出图片名
    file_A = "total_loss.txt"
    file_B = "val_loss.txt"
    A_curve = "Our_loss"
    B_curve = "UNet_loss"
    X_label = "epoch"
    Y_label = "Loss"
    legend = "upper right"
    output_png = "loss2.png"

    # 创建 Plotter 对象，并调用 plot 方法生成图形
    plotter = Plotter(log_directory_A, log_directory_B, output_directory_C, file_A, file_B, A_curve, B_curve, X_label, Y_label, legend, output_png)
    plotter.plot()
